import numpy as np
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import LSTM, Dropout, Dense
from tensorflow.keras.utils import to_categorical

class LSTMModel:
    def __init__(self, input_shape=(1000,1), num_classes=2):
        self.model = Sequential([
            LSTM(64, return_sequences=True, input_shape=input_shape),
            Dropout(0.5),
            LSTM(32),
            Dense(num_classes, activation='softmax')
        ])
        self.model.compile(optimizer='adam', loss='categorical_crossentropy', metrics=['accuracy'])

    def train(self, X, y, epochs=10, batch_size=32):
        y_cat = to_categorical(y)
        self.model.fit(X, y_cat, epochs=epochs, batch_size=batch_size, verbose=0)

    def predict(self, X):
        preds = self.model.predict(X)
        return np.argmax(preds, axis=1)

